

package org.finesys.chat.core.base.service.impl;


import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.chat.core.base.chat.ChatRequest;
import org.finesys.chat.core.base.chat.EmbeddingR;
import org.finesys.chat.core.base.client.ChatClient;
import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.factory.EmbeddingProvider;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.chat.core.base.service.LangEmbeddingService;
import org.finesys.common.core.util.ValidationUtil;
import org.reactivestreams.Publisher;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Service
@RequiredArgsConstructor
public class LangEmbeddingServiceImpl implements LangEmbeddingService {

    private final EmbeddingProvider embeddingProvider;


    @Override
    public Mono<List<EmbeddingR>> embeddingDocs(ChatRequest req) {
        ValidationUtil.ensureNotBlank(req.getKnowledgeId(), "知识库ID");
        //拆分后的文本片段
        List<String> inputs = req.getInput();
        ValidationUtil.ensureNotEmpty(inputs, "输入文本");
        ChatClient chatClient = embeddingProvider.getChatClient(req.getKnowledgeId());
        EmbeddingStore<TextSegment> embeddingStore = embeddingProvider.getEmbeddingStore(req.getKnowledgeId());
        Publisher<EmbeddingResponse> embeddingResponsePublisher = chatClient.embedding(req);
        return Mono.from(embeddingResponsePublisher)
                .flatMap(res -> {
                    //转换为Embedding对象
                    List<Embedding> embeddingVector = res.data().stream().map(r -> {
                        return Embedding.from(r.embedding());
                    }).collect(Collectors.toList());
                    //
                    List<TextSegment> textSegments = inputs.stream().map(TextSegment::from).collect(Collectors.toList());
                    //存储向量
                    List<String> vectorIds = embeddingStore.addAll(embeddingVector, textSegments);
                    //
                    List<EmbeddingR> embeddingRs = new ArrayList<>(vectorIds.size());
                    for (int i = 0; i < vectorIds.size(); i++) {
                        embeddingRs.add(new EmbeddingR()
                                .setVectorId(vectorIds.get(i))
                                .setKnowledgeId(req.getKnowledgeId())
                                .setDocsId(req.getDocId())
                                .setText(textSegments.get(i).getText()));
                    }
                    return Mono.just(embeddingRs);
                });
    }

    @Override
    public Mono<EmbeddingR> embeddingText(ChatRequest req) {
        ValidationUtil.ensureNotBlank(req.getKnowledgeId(), "知识库ID");
        EmbeddingStore<TextSegment> embeddingStore = embeddingProvider.getEmbeddingStore(req.getKnowledgeId());
        ChatClient chatClient = embeddingProvider.getChatClient(req.getKnowledgeId());
        Publisher<EmbeddingResponse> embeddingResponsePublisher = chatClient.embedding(req);

        return Mono.from(embeddingResponsePublisher)
                .flatMap(res -> {
                    List<Double> embeddingVector = res.embedding();
                    Embedding embedding = Embedding.from(embeddingVector);
                    String vectorId = embeddingStore.add(embedding);
                    return Mono.just(new EmbeddingR()
                            .setVectorId(vectorId)
                            .setDocsId(req.getDocId())
                            .setKnowledgeId(req.getKnowledgeId())
                            .setText(req.getInput().toString()));
                });
    }
}
